import cv2
import numpy as np
from PIL import Image
import os

def extract_regions_from_mask(mask_path, image_path, crop_info, resize, threshold=128, padding=10):
    """
    从 mask 中提取感兴趣区域，并映射到原始图片。
    
    Args:
        mask_path (str): mask 图片路径
        image_path (str): 原始图片路径
        crop_info (dict): 裁剪信息，包含 'cropped_size' 和 'crop_box'
        resize (tuple): resize 后的尺寸 (width, height)
        threshold (int): 二值化阈值
        padding (int): 边界扩展大小
    
    Returns:
        list: 包含每个区域信息的字典列表
    """
    # 读取 mask
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    # 二值化
    _, binary_mask = cv2.threshold(mask, threshold, 255, cv2.THRESH_BINARY)
    # 找到连通区域
    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 读取原始图片
    original_image = Image.open(image_path)
    original_array = np.array(original_image)
    
    # 创建保存文件夹
    save_dir = os.path.join('extracted_regions', os.path.basename(image_path).split('.')[0])
    os.makedirs(save_dir, exist_ok=True)
    
    regions = []
    for i, contour in enumerate(contours):
        # 获取边界框
        x, y, w, h = cv2.boundingRect(contour)
        # 添加 padding 并调整为正方形
        x_pad = max(0, x - padding)
        y_pad = max(0, y - padding)
        w_pad = min(mask.shape[1] - x_pad, w + 2 * padding)
        h_pad = min(mask.shape[0] - y_pad, h + 2 * padding)
        size = max(w_pad, h_pad)
        if w_pad < size:
            x_pad = max(0, x_pad - (size - w_pad) // 2)
            w_pad = size
        if h_pad < size:
            y_pad = max(0, y_pad - (size - h_pad) // 2)
            h_pad = size
        
        # 反 resize 到裁剪后图片
        resize_ratio_w = resize[0] / crop_info['cropped_size'][1]  # width
        resize_ratio_h = resize[1] / crop_info['cropped_size'][0]  # height
        x_c = x_pad / resize_ratio_w
        y_c = y_pad / resize_ratio_h
        w_c = w_pad / resize_ratio_w
        h_c = h_pad / resize_ratio_h
        
        # 映射到原始图片
        crop_box = crop_info['crop_box']  # (top, bottom, left, right)
        x_o = int(x_c + crop_box[2])  # left
        y_o = int(y_c + crop_box[0])  # top
        w_o = int(w_c)
        h_o = int(h_c)
        
        # 截取区域
        region = original_array[y_o:y_o+h_o, x_o:x_o+w_o]
        region_image = Image.fromarray(region)
        save_path = os.path.join(save_dir, f'region_{i}.png')
        region_image.save(save_path)
        
        regions.append({
            'save_path': save_path,
            'bbox': (x_o, y_o, x_o+w_o, y_o+h_o)
        })
    
    return regions